# db_handler.py

import mysql.connector.pooling
from contextlib import contextmanager
from mysql.connector import Error
import time
import uuid
import logging
from config import Config
from typing import Optional, List, Dict, Any, Tuple

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class DatabaseHandler:
    def __init__(self, host: str, user: str, password: str, database: str, pool_size: int = 10):
        """初始化数据库连接池
        
        Args:
            host: 数据库主机地址
            user: 数据库用户名
            password: 数据库密码
            database: 数据库名称
            pool_size: 连接池大小
        """
        self.dbconfig = {
            "host": host,
            "user": user,
            "password": password,
            "database": database
        }
        try:
            self.cnxpool = mysql.connector.pooling.MySQLConnectionPool(
                pool_name="mypool",
                pool_size=pool_size,
                **self.dbconfig
            )
        except Error as e:
            logger.error(f"Failed to create connection pool: {e}")
            raise

    @contextmanager
    def get_connection(self, max_retries: int = 10, retry_delay: int = 1):
        """获取数据库连接的上下文管理器
        
        Args:
            max_retries: 最大重试次数
            retry_delay: 重试延迟时间(秒)
            
        Yields:
            Connection: 数据库连接对象
        """
        connection = None
        retries = 0
        while retries <= max_retries:
            try:
                connection = self.cnxpool.get_connection()
                yield connection
                break
            except Error as err:
                if 'pool exhausted' in str(err) and retries < max_retries:
                    logger.warning(f"Connection pool exhausted, retrying ({retries+1}/{max_retries})...")
                    retries += 1
                    time.sleep(retry_delay)
                else:
                    logger.error(f"Database connection error: {err}")
                    raise
            finally:
                if connection and connection.is_connected():
                    connection.close()

    def register_user(self, user_id: str) -> bool:
        """注册新用户
        
        Args:
            user_id: 用户ID
            
        Returns:
            bool: 注册是否成功
        """
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor(prepared=True)
                
                # 检查用户是否已存在
                sql_check = "SELECT 1 FROM sys_user_knowledge WHERE user_id=%s LIMIT 1"
                cursor.execute(sql_check, (user_id,))
                if cursor.fetchone():
                    logger.info(f"User {user_id} already exists")
                    return True
                
                # 创建新用户
                user_uuid = str(uuid.uuid4())
                sql_insert = "INSERT INTO sys_user_knowledge (id, user_id) VALUES (%s, %s)"
                cursor.execute(sql_insert, (user_uuid, user_id))
                connection.commit()
                return True
                
        except Error as e:
            logger.error(f"Failed to register user {user_id}: {e}")
            return False

    def update_knowledge_id(self, user_id: str, knowledge_id: str) -> bool:
        """更新用户的知识库ID
        
        Args:
            user_id: 用户ID
            knowledge_id: 知识库ID
            
        Returns:
            bool: 更新是否成功
        """
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor(prepared=True)
                sql = "UPDATE sys_user_knowledge SET knowledge_id=%s WHERE user_id=%s"
                cursor.execute(sql, (knowledge_id, user_id))
                connection.commit()
                return True
        except Error as e:
            logger.error(f"Failed to update knowledge ID for user {user_id}: {e}")
            return False

    def get_user_api_key(self, user_id: str) -> Optional[str]:
        """获取用户的API密钥
        
        Args:
            user_id: 用户ID
            
        Returns:
            Optional[str]: API密钥，如果不存在则返回None
        """
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor(prepared=True)
                sql = "SELECT api_key FROM sys_user_knowledge WHERE user_id=%s"
                cursor.execute(sql, (user_id,))
                result = cursor.fetchone()
                return result[0] if result and result[0] else None
        except Error as e:
            logger.error(f"Failed to get API key for user {user_id}: {e}")
            return None

    def update_api_key(self, api_key: str, user_id: str) -> bool:
        """更新用户的API密钥
        
        Args:
            api_key: 新的API密钥
            user_id: 用户ID
            
        Returns:
            bool: 更新是否成功
        """
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor(prepared=True)
                sql = "UPDATE sys_user_knowledge SET api_key=%s WHERE user_id=%s"
                cursor.execute(sql, (api_key, user_id))
                connection.commit()
                return True
        except Error as e:
            logger.error(f"Failed to update API key for user {user_id}: {e}")
            return False

    def insert_documents(self, documents: List[Tuple[str, str, str]]) -> bool:
        """批量插入文档记录
        
        Args:
            documents: 文档列表，每个元素为(file_name, user_id, file_ragflow_id)的元组
            
        Returns:
            bool: 插入是否成功
        """
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor(prepared=True)
                sql = """
                    INSERT INTO knowledge_file (file_name, user_id, file_ragflow_id) 
                    VALUES (%s, %s, %s)
                """
                cursor.executemany(sql, documents)
                connection.commit()
                return True
        except Error as e:
            logger.error(f"Failed to insert documents: {e}")
            return False

    def get_user_knowledge_info(self, user_id: str) -> Optional[Dict[str, str]]:
        """获取用户的知识库信息
        
        Args:
            user_id: 用户ID
            
        Returns:
            Optional[Dict[str, str]]: 包含knowledge_id和api_key的字典，如果不存在则返回None
        """
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor(prepared=True)
                sql = """
                    SELECT knowledge_id, api_key 
                    FROM sys_user_knowledge 
                    WHERE user_id=%s 
                    LIMIT 1
                """
                cursor.execute(sql, (user_id,))
                result = cursor.fetchone()
                if result and result[0]:
                    return {
                        "knowledge_id": result[0],
                        "api_key": result[1]
                    }
                return None
        except Error as e:
            logger.error(f"Failed to get knowledge info for user {user_id}: {e}")
            return None

    def get_user_documents(self, user_id: str, filename: Optional[str] = None, 
                          file_status: Optional[int] = None) -> List[Dict[str, Any]]:
        """获取用户的文档列表
        
        Args:
            user_id: 用户ID
            filename: 可选的文件名过滤
            file_status: 可选的文件状态过滤
            
        Returns:
            List[Dict[str, Any]]: 文档信息列表
        """
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor(prepared=True)
                sql = """
                    SELECT id, file_name, document_info, file_status, addtime 
                    FROM knowledge_file 
                    WHERE user_id=%s
                """
                params = [user_id]

                if filename:
                    sql += " AND file_name=%s"
                    params.append(filename)

                if file_status is not None:
                    sql += " AND file_status=%s"
                    params.append(file_status)

                cursor.execute(sql, tuple(params))
                results = cursor.fetchall()
                
                return [{
                    "document_id": row[0],
                    "file_name": row[1],
                    "document_info": row[2],
                    "file_status": row[3],
                    "addtime": row[4].strftime("%Y-%m-%d %H:%M:%S") if row[4] else None
                } for row in results]
                
        except Error as e:
            logger.error(f"Failed to get documents for user {user_id}: {e}")
            return []

    def update_knowledge_base_id(self, file_id: str, knowledge_base_id: str) -> bool:
        """
        更新文件的知识库ID
        
        Args:
            file_id: 文件ID
            knowledge_base_id: 知识库ID
            
        Returns:
            bool: 是否更新成功
        """
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor(prepared=True)
                query = "UPDATE knowledge_file SET knowledge_base_id = %s WHERE id = %s"
                cursor.execute(query, (knowledge_base_id, file_id))
                connection.commit()
                return True
        except Exception as e:
            logger.error(f"更新 knowledge_base_id 失败: {str(e)}")
            return False

    def test_connection(self) -> bool:
        """测试数据库连接是否正常
        
        Returns:
            bool: 连接是否成功
        """
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor()
                cursor.execute("SELECT 1")
                result = cursor.fetchone()
                return result and result[0] == 1
        except Exception as e:
            logger.error(f"Database connection test failed: {e}")
            return False

# 初始化数据库处理器实例
db_handler = DatabaseHandler(host='db.cobona.cn',
                             user='root',
                             password='kebona',
                             database='skxy')
